#include	<sys/time.h>
#include	<algorithm>
#include	<iostream>
#include	<unistd.h>
#include	"utf16.h"
#include	<cstring>
#include	<sstream>
#include	<cfloat>
#include	<cstdio>
#include	<vector>
#include	<cmath>
#include	<ctime>
#include	<omp.h>
using	namespace	std;
#define	_F2C	9.48683298050514f
#define	_C2F	0.105409255338946f
uint64_t	prng=time(NULL);
static inline uint64_t wyrand(uint64_t	*seed) {
    *seed+=0xa0761d6478bd642full;
    uint64_t  see1=*seed^0xe7037ed1a0b428dbull;
    see1*=(see1>>32)|(see1<<32);
    return	(*seed*((*seed>>32)|(*seed<<32)))^((see1>>32)|(see1<<32));
}
static	inline float wy2gau(uint64_t r) {
    const float _wynorm=1.0f/(1ull<<20);
    return ((r&0x1fffff)+((r>>21)&0x1fffff)+((r>>42)&0x1fffff))*_wynorm-3.0f;
}
static inline double wy2u01(uint64_t r) {
    const double _wynorm=1.0/(1ull<<52);
    return (r>>12)*_wynorm;
}
inline void wymum(uint64_t *A,	uint64_t *B) {
    uint64_t	hh=(*A>>32)*(*B>>32), hl=(*A>>32)*(uint32_t)*B, lh=(uint32_t)*A*(*B>>32), ll=(uint64_t)(uint32_t)*A*(uint32_t)*B;
    *A=((hl>>32)|(hl<<32))^hh;
    *B=((lh>>32)|(lh<<32))^ll;
}
inline	uint64_t	wyhash64(uint64_t	A,	uint64_t	B) {
    A^=0xa0761d6478bd642full;
    B^=0xe7037ed1a0b428dbull;
    wymum(&A,&B);
    A^=0xe7037ed1a0b428dbull;
    B^=0xa0761d6478bd642full;
    wymum(&A,&B);
    return	A^B;
}

template<class	T,	uint32_t	N>
struct	tensor {
    T	*p;
    tensor() {
        p=(T*)aligned_alloc(64,N*sizeof(T));
    }
//	tensor(){	p=(T*)_aligned_malloc(N*sizeof(T),64);	}
    ~tensor() {
        free(p);
    }
//	~tensor(){	_aligned_free(p);	}
    bool	load(FILE	*F) {
        return	fread(p,N*sizeof(T),1,F)==1;
    }
};

template<uint32_t	N>
float	f2c(float	*A,	int8_t	*B) {
    float	ma=0;
    for(uint32_t	i=0;	i<N;	i++)	ma=fmaxf(ma,fabsf(A[i]));
    ma=90/ma;
    for(uint32_t	i=0;	i<N;	i++)	B[i]=roundf(A[i]*ma);
    return	1.0f/ma;
}

template<uint32_t	N>
float	f2cn(float	*A,	int8_t	*B) {
    for(uint32_t	i=0;	i<N;	i++)	B[i]=roundf(A[i]*_F2C);
    return	_C2F;
}

template<uint32_t	N>
int32_t	dot(int8_t	*A,	int8_t	*B) {
    int32_t	s=0;
    for(uint32_t	i=0;	i<N;	i++)	s+=(int32_t)A[i]*(int32_t)B[i];
    return	s;
}

template<uint32_t	R0,	uint32_t	R1>
struct	linear {
    tensor<int8_t,R0>	inc;
    tensor<int8_t,R0*R1>	wei;
    tensor<float,R1>	out;
    void	load(FILE	*F) {
        wei.load(F);
    }
    void	fw(tensor<float,R0>	&inp) {
        float	alf=_C2F*f2c<R0>(inp.p,inc.p)/sqrtf(R0);
        for(uint32_t	i=0;	i<R1;	i++) out.p[i]=alf*dot<R0>(inc.p,wei.p+i*R0);
    }
};
template<uint32_t	R,	uint32_t	H>
void	rmsnorm(tensor<float,R>	&inp) {
    uint32_t	r=R/H;
    for(uint32_t	h=0;	h<H;	h++) {
        float	nor=0,	*p=inp.p+h*r;
        for(uint32_t	i=0;	i<r;	i++)	nor+=p[i]*p[i];
        nor=sqrtf(r/nor);
        for(uint32_t	i=0;	i<r;	i++)	p[i]*=nor;
    }
}
template<uint32_t	R>
void	softmax(tensor<float,R>	&inp) {
    float	ma=-FLT_MAX,	sum=0;
    for(uint32_t	i=0;	i<R;	i++)	if(inp.p[i]>ma)	ma=inp.p[i];
    for(uint32_t	i=0;	i<R;	i++)	sum+=(inp.p[i]=expf(inp.p[i]-ma));
    for(uint32_t	i=0;	i<R;	i++)	inp.p[i]/=sum;
}
template<uint32_t	R,	uint32_t	C,	uint32_t	H,	uint32_t	S>
struct	transformer {
    tensor<int8_t,H*S*C>	vk;
    tensor<int8_t,C*H*S>	pn;
    tensor<int8_t,H*C>	p;
    tensor<int8_t,S>	cq;
    tensor<int8_t,C>	ca;
    linear<R,6*H*S>	x;
    linear<2*H*S,R>	o;
    tensor<float,2*H*S>	tmp;
    tensor<float,H*S>	va;
    tensor<float,R>	&out=o.out;
    void	load(FILE	*F) {
        p.load(F);
        x.load(F);
        o.load(F);
    }
    void	fw(tensor<float,R>	&inp,	uint32_t	col,	uint32_t	para) {
        x.fw(inp);
        rmsnorm<6*H*S,6*H>(x.out);
        f2cn<H*S>(x.out.p,vk.p+col*H*S);
        for(size_t	i=0;	i<H*S;	i++)	pn.p[i*C+col]=roundf(x.out.p[2*H*S+i]*_F2C);
        for(uint32_t	h=0;	h<H;	h++) {
            float	alf=M_SQRT2/sqrtf(S)*_C2F*_C2F,	a[C]= {};
            f2cn<S>(x.out.p+H*S+h*S,cq.p);
            for(uint32_t	j=para;	j<C;	j++) {
                uint32_t	i=(j+1+col)%C;
                a[i]=expf(dot<S>(vk.p+i*H*S+h*S,cq.p)*alf+p.p[h*C+C-1-j]*_C2F);
            }
            float	s=_C2F*f2c<C>(a,ca.p);
            for(uint32_t	i=0;	i<S;	i++)	va.p[h*S+i]=s*dot<C>(pn.p+(h*S+i)*C,ca.p);
        }
        for(uint32_t	t=0;	t<2;	t++) {
            float	*pu=x.out.p+(3+t)*H*S,	*pv=t?x.out.p+(4+t)*H*S:va.p,	*po=tmp.p+t*H*S;
            for(uint32_t	i=0;	i<H*S;	i++)	po[i]=pu[i]*pv[i];
        }
        rmsnorm<2*H*S,2*H>(tmp);
        o.fw(tmp);
        for(uint32_t	i=0;	i<R;	i++)	o.out.p[i]+=inp.p[i];
    }
};
template<uint32_t	C,	uint32_t	E,	uint32_t	D,	uint32_t	H,	uint32_t	S,	uint32_t	O>
struct	HomoSapiens {
    uint32_t	curr=0,	para=C-1;
    uint16_t	code[O],table[65536];
    vector<uint16_t>	his;
    tensor<float,E>	emb;
    transformer<E,C,H,S>	tra[D];
    linear<E,O>	out;
    float	vs[O];

    bool	load(const	char	*F) {
        FILE* f=fopen(F, "rb");
        if(f==NULL)	return	false;
        uint32_t	x;
        if(fread(&x,4,1,f)!=1||x!=C)	fprintf(stderr,"C=%u\n",x);
        if(fread(&x,4,1,f)!=1||x!=E)	fprintf(stderr,"E=%u\n",x);
        if(fread(&x,4,1,f)!=1||x!=D)	fprintf(stderr,"D=%u\n",x);
        if(fread(&x,4,1,f)!=1||x!=H)	fprintf(stderr,"H=%u\n",x);
        if(fread(&x,4,1,f)!=1||x!=S)	fprintf(stderr,"H=%u\n",x);
        if(fread(&x,4,1,f)!=1||x!=O)	fprintf(stderr,"O=%u\n",x);
        if(fread(code,O*2,1,f)!=1)	return	false;
        for(uint32_t i=0; i<D; i++)	tra[i].load(f);
        out.load(f);
        fclose(f);
        for(size_t	i=0;	i<65536;	i++)	table[i]=0;
        for(size_t	i=0;	i<O;	i++)	table[code[i]]=i;
        return	true;
    }
    void	reset(void) {
        curr=0;
        para=C-1;
        his.clear();
    }
    uint16_t	sample(uint16_t	x) {
        his.push_back(x);
        memset(vs,0,sizeof(float)*O);
        for(uint32_t	k=his.size()<C?0:his.size()-C;	k<his.size();	k++) {
            uint32_t	l=1;
            while(l<=k&&his[k-l]==his[his.size()-l])	l++;
            vs[table[his[k]]]+=(expf(l)-1)/(his.size()-k);
        }
        for(uint32_t	r=0;	r<E;	r++)	emb.p[r]=wy2gau(wyhash64(x,r));
        for(uint32_t	d=0;	d<D;	d++)	tra[d].fw(d?tra[d-1].out:emb,curr,para);
        rmsnorm<E,1>(tra[D-1].out);
        out.fw(tra[D-1].out);
        for(uint32_t    i=0;    i<O;    i++)	out.out.p[i]=code[i]=='#'?-1000:2*(out.out.p[i]-vs[i]);
        softmax<O>(out.out);
//        double	sum=0;
//        for(uint32_t    i=0;    i<O;    i++)	sum+=(out.out.p[i]=fmaxf(out.out.p[i]-1.0f/O,0));
        double  ran=wy2u01(wyrand(&prng)),  sum1=0;
        uint16_t        ret=0;
        for(size_t      i=0;    i<O;    i++) {
            sum1+=out.out.p[i];
            if(sum1>=ran) {
                ret=i;
                break;
            }
        }
        curr=(curr+1)%C;
        if(para)	para--;
        return	code[ret];
    }
    string	generate(string	inp) {
        uint32_t	t=0;
        vector<uint16_t>	v,u;
        uint16_t	c;
        v.resize(inp.size());
        v.resize(utf8_to_utf16((uint8_t*)inp.data(),inp.size(),v.data(),v.size()));
        for(uint32_t	i=0;	i<v.size();	i++)	c=sample(v[i]);
        while(c!='\n') {
            u.push_back(c);
            c=sample(c);
        }
        u.push_back(c);
        c=sample(c);
        string	ret;
        ret.resize(u.size()*4);
        ret.resize(utf16_to_utf8(u.data(),u.size(),(uint8_t*)ret.data(),ret.size()));
        return	ret;
    }
};
#include	"config"
int	main(int	ac,	char	**av) {
    HomoSapiens<context,embed,depth,headn,heads,voca>	model;
    string	model_file="model";
    int	opt;
    while((opt=getopt(ac,	av,	"m:"))>=0) {
        switch(opt) {
        case	'm':
            model_file=optarg;
            break;
        }
    }
    if(!model.load(model_file.c_str())) {
        fprintf(stderr,"fail to load %s\n",model_file.c_str());
        return	0;
    }
    string	s;
    for(;;) {
        getline(cin,s);
        cout<<model.generate(s+"\n");
    }
    return	0;
}

